"""
Description: 五子棋 v1 版本，实现了基础的落子功能，有一方获胜后可以重新开始
Version: 1.0
Author: marshal
Date: 2024-06-02 0:52
LastEditors: marshal
LastEditTime: 2024-06-02 0:52
"""

import pygame as pg
import sys
from pygame.locals import *
import numpy as np

pg.init()
# 颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)

# 棕色背景颜色
BOARD_COLOR = (205, 133, 63)
LINE_COLOR = BLACK
# 棋盘参数
MARGINS = 50  # 边距
SPACING = 40  # 格子间距
LINES = 15  # 线条数
# 屏幕尺寸
SCREEN_SIZE = MARGINS * 2 + (LINES - 1) * SPACING
CHESS_RADIUS = 18  # 棋子半径
# 棋子状态
EMPTY = 0
BLACK_CHESS = 1
WHITE_CHESS = 2
# 记录棋子状态，数据类型是二维数组
CHESS_STATUS = np.zeros([15, 15], dtype=int)
# 标记点位
MARK_POINTS = [(3, 3), (3, 11), (7, 7), (11, 3), (11, 11)]

logo = pg.image.load("../pictures/logo.png")
pg.display.set_icon(logo)
pg.display.set_caption("五子棋")
screen = pg.display.set_mode((SCREEN_SIZE, SCREEN_SIZE))


def convert_board_pos(x, y):
    """
    坐标转换：将 Pygame 屏幕的坐标转换为棋盘坐标
    :param x:
    :param y:
    :return:
    """
    if check_pos(x, y):
        row = round((x - MARGINS) / SPACING)
        col = round((y - MARGINS) / SPACING)
        return row, col


def convert_pg_pos(x, y):
    return MARGINS + SPACING * x, MARGINS + SPACING * y


def draw_board(margins, spacing):
    """
    绘制棋盘
    :param margins: 棋盘边距
    :param spacing: 每个格子的距离
    :return:
    """
    # 先画外侧线
    pg.draw.lines(screen, LINE_COLOR, True,
                  [(margins, margins), (margins, SCREEN_SIZE - margins), (SCREEN_SIZE - margins, SCREEN_SIZE - margins),
                   (SCREEN_SIZE - margins, margins)], 3)

    # 画内侧线
    for i in range(1, LINES - 1):
        # 画竖线
        pg.draw.line(screen, LINE_COLOR, (margins + spacing * i, margins),
                     (margins + spacing * i, SCREEN_SIZE - margins))
        # 画横线
        pg.draw.line(screen, LINE_COLOR, (margins, margins + spacing * i),
                     (SCREEN_SIZE - margins, margins + spacing * i))
    # 标记特殊点位
    for points in MARK_POINTS:
        pos_x = points[0]
        pos_y = points[1]
        radius = 6 if pos_x == int(LINES / 2) else 4
        pg.draw.circle(screen, LINE_COLOR, (margins + spacing * pos_x, margins + spacing * pos_y), radius)


def check_pos(x, y):
    """
    判断坐标是否在棋盘内
    :param x: 横坐标
    :param y: 纵坐标
    :return:
    """
    return MARGINS <= x <= SCREEN_SIZE - MARGINS and MARGINS <= y <= SCREEN_SIZE


def play(x, y, color):
    """
    下棋方法，实现思路是记录鼠标左键点击的坐标的颜色
    :return:
    """
    if is_win()[0] != EMPTY:
        return False
    if check_pos(x, y):
        row, col = convert_board_pos(x, y)
        # 如果该点位的状态是 EMPTY ，说明该位置为没有落子，则可以设置颜色
        if CHESS_STATUS[row][col] == EMPTY:
            CHESS_STATUS[row][col] = BLACK_CHESS if color == BLACK else WHITE_CHESS
            return True
    return False


def is_win():
    """
    判断斜边方向上是否存在五子连珠
    :return:
    """
    # 1. 判断水平方向
    for i in range(LINES):
        pos1 = []
        pos2 = []
        for j in range(LINES):
            if CHESS_STATUS[i][j] == BLACK_CHESS:
                pos1.append((i, j))
            else:
                pos1 = []
            if len(pos1) >= 5:
                return BLACK_CHESS, pos1
            if CHESS_STATUS[i][j] == WHITE_CHESS:
                pos2.append((i, j))
            else:
                pos2 = []
            if len(pos2) >= 5:
                return WHITE_CHESS, pos2
    # 2. 判断竖直方向
    for j in range(LINES):
        pos1 = []
        pos2 = []
        for i in range(LINES):
            if CHESS_STATUS[i][j] == BLACK_CHESS:
                pos1.append((i, j))
                if len(pos1) >= 5:
                    return BLACK_CHESS, pos1
            else:
                pos1 = []
            if CHESS_STATUS[i][j] == WHITE_CHESS:
                pos2.append((i, j))
                if len(pos2) >= 5:
                    return WHITE_CHESS, pos2
            else:
                pos2 = []
    # 3. 正斜方向五子连珠
    for i in range(LINES):
        for j in range(LINES):
            pos1 = []
            pos2 = []
            for k in range(LINES):
                if i + k >= 15 or j + k >= 15:
                    break
                if CHESS_STATUS[i + k][j + k] == BLACK_CHESS:
                    pos1.append((i + k, j + k))
                else:
                    pos1 = []
                if len(pos1) >= 5:
                    return BLACK_CHESS, pos1
                if CHESS_STATUS[i + k][j + k] == WHITE_CHESS:
                    pos2.append((i + k, j + k))
                else:
                    pos2 = []
                if len(pos2) >= 5:
                    return WHITE_CHESS, pos2
    # 4. 反斜方向五子连珠
    for i in range(LINES):
        for j in range(LINES):
            pos1 = []
            pos2 = []
            for k in range(LINES):
                if i + k == 15 or j - k < 0:
                    break
                if CHESS_STATUS[i + k][j - k] == BLACK_CHESS:
                    pos1.append((i + k, j - k))
                else:
                    pos1 = []
                if len(pos1) >= 5:
                    return BLACK_CHESS, pos1
                if CHESS_STATUS[i + k][j - k] == WHITE_CHESS:
                    pos2.append((i + k, j - k))
                else:
                    pos2 = []
                if len(pos2) >= 5:
                    return WHITE_CHESS, pos2

    return EMPTY, []


def reset():
    """
    重置方法
    :return:
    """
    global CHESS_STATUS
    screen.fill(BOARD_COLOR)
    draw_board(MARGINS, SPACING)
    CHESS_STATUS = np.zeros([15, 15], dtype=int)
    pg.display.update()


def draw_chess():
    """
    绘制棋子
    :return:
    """
    # 如果存在非 0 的数值则说明不是空棋盘
    if np.any(CHESS_STATUS > 0):
        for row in range(LINES):
            for col in range(LINES):
                status = CHESS_STATUS[row][col]
                if status != EMPTY:
                    pos = (MARGINS + SPACING * row, MARGINS + SPACING * col)
                    pg.draw.circle(screen, BLACK if status == BLACK_CHESS else WHITE, pos, CHESS_RADIUS)


def draw_chess_box(x, y):
    """
    绘制落子框，当鼠标悬停在棋盘上时，显示是否可以落子
    :param x: 鼠标横坐标
    :param y: 鼠标纵坐标
    :return:
    """
    if check_pos(x, y):
        row, col = convert_board_pos(x, y)
        pos = [
            (MARGINS + SPACING * row - CHESS_RADIUS, MARGINS + SPACING * col - CHESS_RADIUS),
            (MARGINS + SPACING * row + CHESS_RADIUS, MARGINS + SPACING * col - CHESS_RADIUS),
            (MARGINS + SPACING * row + CHESS_RADIUS, MARGINS + SPACING * col + CHESS_RADIUS),
            (MARGINS + SPACING * row - CHESS_RADIUS, MARGINS + SPACING * col + CHESS_RADIUS)
        ]
        pg.draw.lines(screen, RED, True, pos, 2)


def main():
    running = True
    is_black = True
    font = pg.font.Font("../fonts/FangZhengHeiTiJianTi-1.ttf", 20)
    while running:
        # 检测退出事件
        for event in pg.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == MOUSEBUTTONDOWN and event.button == 1:
                x, y = pg.mouse.get_pos()
                # 落子
                if play(x, y, BLACK if is_black else WHITE):
                    is_black = not is_black
        # 绘制棋盘
        screen.fill(BOARD_COLOR)
        draw_board(MARGINS, SPACING)

        # 拿到鼠标位置，绘制落子框
        x, y = pg.mouse.get_pos()
        draw_chess_box(x, y)
        # 绘制棋子
        draw_chess()
        # 检查是否赢了
        status, pos = is_win()
        if status != EMPTY:
            if status == BLACK_CHESS:
                text = font.render("黑棋赢了，将在5秒后重新开始", True, BLACK)
            else:
                text = font.render("白棋赢了，将在5秒后重新开始", True, WHITE)
            text_pos = text.get_rect(centerx=SCREEN_SIZE / 2, y=15)
            screen.blit(text, text_pos)
            pg.display.update()
            pg.time.delay(5000)
            reset()
        else:
            if is_black:
                text = font.render("黑棋回合", True, BLACK)
            else:
                text = font.render("白棋回合", True, WHITE)
            text_pos = text.get_rect(centerx=SCREEN_SIZE / 2, y=15)
            screen.blit(text, text_pos)
        pg.display.update()
    pg.quit()
    sys.exit()


if __name__ == '__main__':
    main()
